#!/usr/bin/env python
# -*- coding: gb2312 -*-

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

# Mexican hat wavelet function
def mexican_hat(t):
    return (1 - t**2) * np.exp(-t**2 / 2)

# Create a test signal with multiple features
def create_signal():
    t = np.linspace(-6, 6, 300)
    signal = np.zeros_like(t)
    # Create multiple pulses with different widths
    signal += 2.0 * np.exp(-10 * (t + 4)**2)  # Narrow pulse
    signal += 1.5 * np.exp(-3 * (t + 1)**2)   # Medium pulse
    signal += 1.8 * np.exp(-5 * (t - 2)**2)   # Sharp pulse
    signal += 1.2 * np.exp(-2 * (t - 5)**2)   # Wide pulse
    return t, signal

# Create enhanced visualization
def create_enhanced_plots():
    # Generate signal and wavelet
    t, signal = create_signal()
    dt = t[1] - t[0]
    
    t_wavelet = np.linspace(-3, 3, 100)
    wavelet = mexican_hat(t_wavelet)
    
    # Calculate convolution
    conv_result = np.convolve(signal, wavelet, mode='same') * dt
    
    # Create a figure with 4 subplots
    plt.figure(figsize=(14, 10), dpi=120)
    
    # Subplot 1: Original signal
    plt.subplot(4, 1, 1)
    plt.plot(t, signal, 'b-', linewidth=2)
    plt.title('1. Original Signal with Multiple Pulses', fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.ylabel('Amplitude')
    
    # Subplot 2: Mexican hat wavelet
    plt.subplot(4, 1, 2)
    plt.plot(t_wavelet, wavelet, 'r-', linewidth=2)
    plt.fill_between(t_wavelet, wavelet, alpha=0.2, color='red')
    plt.title('2. Mexican Hat Wavelet', fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.ylabel('Amplitude')
    # Add formula annotation
    plt.text(-2.8, 0.8, r'$\psi(t) = (1-t^2)e^{-t^2/2}$', fontsize=12)
    
    # Subplot 3: Convolution result
    plt.subplot(4, 1, 3)
    plt.plot(t, conv_result, 'g-', linewidth=2)
    plt.title('3. Convolution Result', fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.ylabel('Convolution Value')
    
    # Subplot 4: Signal and convolution overlay
    plt.subplot(4, 1, 4)
    plt.plot(t, signal, 'b-', linewidth=1.5, alpha=0.7, label='Signal')
    plt.plot(t, conv_result/5, 'g-', linewidth=2, label='Normalized Conv')
    plt.title('4. Signal and Normalized Convolution', fontsize=14)
    plt.grid(True, alpha=0.3)
    plt.xlabel('Time')
    plt.ylabel('Amplitude')
    plt.legend(loc='best')
    
    plt.tight_layout()
    plt.savefig('enhanced_wavelet_analysis.png', dpi=120, bbox_inches='tight')
    plt.close()
    
    # Create a separate figure showing convolution at key positions
    plt.figure(figsize=(12, 8), dpi=100)
    
    # Find positions of the peaks in the original signal
    from scipy.signal import find_peaks
    peaks, _ = find_peaks(signal, height=0.5)
    
    # Select up to 3 key positions to show
    key_positions = peaks[:min(3, len(peaks))]
    
    for i, pos_idx in enumerate(key_positions):
        pos_time = t[pos_idx]
        
        plt.subplot(len(key_positions), 1, i+1)
        
        # Plot the signal
        plt.plot(t, signal, 'b-', linewidth=1.5, label='Signal')
        
        # Create wavelet at current position
        wavelet_t = np.linspace(pos_time - 3, pos_time + 3, 100)
        wavelet_at_pos = mexican_hat(wavelet_t - pos_time)
        
        # Plot the wavelet at current position (scaled for visibility)
        plt.plot(wavelet_t, wavelet_at_pos * 0.5, 'r-', linewidth=1.5, label='Wavelet')
        
        # Calculate overlapping region for product
        overlap_mask = (wavelet_t >= t[0]) & (wavelet_t <= t[-1])
        wavelet_overlap = wavelet_at_pos[overlap_mask]
        t_overlap = wavelet_t[overlap_mask]
        
        # Find corresponding indices in the signal
        idx_start = np.searchsorted(t, t_overlap[0])
        idx_end = np.searchsorted(t, t_overlap[-1]) + 1
        idx_end = min(idx_end, len(t))
        
        # Create signal segment for overlap region
        t_segment = t[idx_start:idx_end]
        signal_segment = signal[idx_start:idx_end]
        
        # Calculate product (truncate if needed)
        min_len = min(len(signal_segment), len(wavelet_overlap))
        product = signal_segment[:min_len] * wavelet_overlap[:min_len]
        
        # Plot product
        plt.fill_between(t_segment[:min_len], product, alpha=0.3, color='green', label='Product')
        
        plt.title(f'Convolution at t={pos_time:.2f}', fontsize=12)
        plt.grid(True, alpha=0.3)
        plt.legend(loc='best')
    
    plt.tight_layout()
    plt.savefig('convolution_positions.png', dpi=100, bbox_inches='tight')
    plt.close()
    
    print("\nGenerated enhanced visualization files:")
    print("1. enhanced_wavelet_analysis.png - Comprehensive wavelet analysis")
    print("2. convolution_positions.png - Convolution at key signal positions")

def main():
    print("Running enhanced Mexican hat wavelet demonstration with GB2312 encoding...")
    create_enhanced_plots()
    print("\nEnhanced demonstration completed successfully!")

if __name__ == "__main__":
    main()